#pragma once

namespace infer {

/**
 * @brief Buffer枚举类型, 选择不同类型Buffer
 */
enum class BufferType {
  Host,   // Host 内存
  Device, // Device 内存
  Shared, // 共享内存
};

enum class DataType {
  INT8,  // int8
  INT32, // int32
  INT64, // int64
  FP16,  // float16
  FP32,  // float32
  BF16,  // bfloat16
};

struct Tensor {
public:
  Tensor(const std::string name, const std::vector<int64_t>& shape,
         const DataType dtype, bool is_input) {
    this->name = name;
    this->shape = shape;
    this->dtype_ = dtype;
    this->is_input = is_input;
  }

  std::string name;           // Tensor 名称
  std::vector<int64_t> shape; // Tensor 形状
  bool is_input;              // 是否为输入 Tensor

  void* hostPtr() { return nullptr; }   // TODO 获取 Host 内存指针
  void* devicePtr() { return nullptr; } // TODO 获取 Device 内存指针

  void toHost() {}   // TODO 将 Tensor 数据从 Device 拷贝到 Host
  void toDevice() {} // TODO 将 Tensor 数据从 Host 拷贝到 Device

private:
  // DataType 待定义
  DataType dtype_ = DataType::FP16; // 数据类型
  size_t bytes_;   // Tensor 大小(字节数)
};

} // namespace infer